#ifndef XLOGY_N_D_H
#define XLOGY_N_D_H

#include "kernel_operator.h"

namespace Xlogy {
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 1;
constexpr int32_t BYTE_BLOCK = 32;
constexpr int32_t BYTE_REPEAT = 256;

constexpr int32_t TMP_BUF_COUNT = 1;

constexpr float ZERO = 0.0;
constexpr float POS_ONE = 1.0;

template <typename T>
class XlogyND {
public:
    __aicore__ inline XlogyND() = default;
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, const XlogyTilingData* __restrict tilingData);
    __aicore__ inline void Process();

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyIn(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void Compute(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void CopyOut(int64_t gmOffset, int64_t dataCount);

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> x1Que, x2Que;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQue;
    TBuf<QuePosition::VECCALC> tempValBuf;
    GlobalTensor<T> x1GM, x2GM, yGM;
    LocalTensor<T> tempValLT;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;
    uint64_t perRepeatCount = 0;

    // tiling params
    uint64_t ubMaxProcCount = 0;
    uint64_t totalDataCount = 0;
    uint64_t loopCount = 0;
    uint64_t tailCount = 0;
};

template <typename T>
__aicore__ inline void XlogyND<T>::Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                        const XlogyTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);
    perRepeatCount = BYTE_REPEAT / sizeof(T);
    ubMaxProcCount = tilingData->ubMaxProcCount;
    totalDataCount = tilingData->totalDataCount;
    loopCount = tilingData->loopCount;
    tailCount = tilingData->tailCount;

    x1GM.SetGlobalBuffer((__gm__ T*)x1, totalDataCount);
    x2GM.SetGlobalBuffer((__gm__ T*)x2, totalDataCount);
    yGM.SetGlobalBuffer((__gm__ T*)y, totalDataCount);

    uint64_t singleBufferSize = ubMaxProcCount * sizeof(T);
    pipe.InitBuffer(x1Que, BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(x2Que, BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(yQue, BUFFER_NUM, singleBufferSize);

    pipe.InitBuffer(tempValBuf, singleBufferSize * TMP_BUF_COUNT);

    tempValLT = tempValBuf.Get<T>();
}

template <typename T>
__aicore__ inline void XlogyND<T>::Process() {
    int64_t gmOffset = 0;
    for (int64_t i = 0; i < loopCount; i++) {
        CopyIn(gmOffset, ubMaxProcCount);
        Compute(gmOffset, ubMaxProcCount);
        CopyOut(gmOffset, ubMaxProcCount);
        gmOffset += ubMaxProcCount;
    }
    if (tailCount) {
        int64_t alignCopyCount = CeilAlignA2B(tailCount, perBlockCount);
        int64_t alignComputeCount = CeilAlignA2B(tailCount, perRepeatCount);
        CopyIn(gmOffset, alignCopyCount);
        Compute(gmOffset, alignComputeCount);
        CopyOut(gmOffset, alignCopyCount);
    }
}

template <typename T>
__aicore__ inline void XlogyND<T>::CopyIn(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> x1InLT = x1Que.AllocTensor<T>();
    LocalTensor<T> x2InLT = x2Que.AllocTensor<T>();
    DataCopy(x1InLT, x1GM[gmOffset], dataCount);
    DataCopy(x2InLT, x2GM[gmOffset], dataCount);
    x1Que.EnQue(x1InLT);
    x2Que.EnQue(x2InLT);
}

template <typename T>
__aicore__ inline void XlogyND<T>::Compute(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> x1LT = x1Que.DeQue<T>();
    LocalTensor<T> x2LT = x2Que.DeQue<T>();
    LocalTensor<T> yLT = yQue.AllocTensor<T>();

    LocalTensor<T> t1 = tempValLT;
    LocalTensor<uint8_t> mask = tempValLT.template ReinterpretCast<uint8_t>();
    Duplicate(t1, T(ZERO), dataCount);
    Compare(mask, x1LT, t1, CMPMODE::NE, dataCount);
    Select(x2LT, mask, x2LT, T(POS_ONE), SELMODE::VSEL_TENSOR_SCALAR_MODE, dataCount);

    Ln(x2LT, x2LT, dataCount);
    Mul(yLT, x1LT, x2LT, dataCount);

    yQue.EnQue(yLT);
    x1Que.FreeTensor(x1LT);
    x2Que.FreeTensor(x2LT);
}

template <typename T>
__aicore__ inline void XlogyND<T>::CopyOut(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> yOutLT = yQue.DeQue<T>();
    DataCopy(yGM[gmOffset], yOutLT, dataCount);
    yQue.FreeTensor(yOutLT);
}
}  // namespace Xlogy

#endif  // XLOGY_N_D_H